package com.agent;

import javassist.*;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Arrays;

/**
 * 使用javaassit修改字节码文件
 * <p>
 * Javassist (JAVA programming ASSISTant) 是在 Java 中编辑字节码的类库;
 * 它使 Java 程序能够在运行时定义一个新类, 并在 JVM 加载时修改类文件。
 * 其他的工具有ASM、ByteBuddy.
 * <p>
 * Javassist主要类有ClassPool、CtClass、CtMethod，其中，ClassPool 是一个存储 CtClass 的 Hash 表，
 * 类的名称作为 Hash 表的 key。ClassPool 的 get() 函数用于从 Hash 表中查找 key 对应的 CtClass 对象。
 * 如果没有找到，get() 函数会创建并返回一个新的 CtClass 对象，这个新对象会保存在 Hash 表中
 * 类 CtClass 是类文件的抽象表示形式
 * 程序需要首先获取一个 ClassPool 对象，才能通过Javassist 控制字节码修改。
 *
 * @author dingrz
 */
public class MonitorDemoTransformer implements ClassFileTransformer {

    /**
     * 此demo只对com.example.springbootdemo.controller 包下的类做字节码转换
     * 处理 className = com/example/springbootdemo/controller/HelloController
     */
    private String basePackageName = "com.example.springbootdemo.controller";

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

        className = className.replace("/", ".");

        if (!className.startsWith(basePackageName)) {
            return null;
        }
        ClassPool pool = ClassPool.getDefault();
        CtClass ctClass = null;
        try {
            ctClass = pool.get(className);
            System.out.println("ctClass.getClass().getName() = " + ctClass.getClass().getName());
            CtMethod[] declaredMethods = ctClass.getDeclaredMethods();//获取类中声明的所有方法。
            System.out.println("当前要增强的类："+className+"\n本类中声明方法如下：");
            Arrays.asList(declaredMethods).forEach(m -> System.out.println(m.getName()));
            for (CtBehavior method : declaredMethods) {
                if (method.getName().startsWith("say")) {
                    System.out.println("当前要增强的方法：" + method.getName());
                    method.addLocalVariable("start", CtClass.longType);
                    method.insertBefore("start = System.currentTimeMillis();\n");
                    method.insertAfter("System.out.println(\" 执行方法[" + method.getName() + "]耗时:[\" +(System.currentTimeMillis()-start)+ \"ms]\");");
                }
            }
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }
}
